from monai.networks.nets.basic_unet import BasicUNet
import torch.nn as nn 
import torch 

class UNet3D(nn.Module):
    def __init__(self, in_ch, out_ch) -> None:
        super().__init__()
        self.model = BasicUNet(spatial_dims=3, 
                                in_channels=in_ch, 
                                out_channels=out_ch,
                                features=[16, 16, 32, 64, 128, 16])

    
    def forward(self, x):
        return self.model(x)
        


